package xyz.thoughtset.viewer.ai.mcp.server.config;

import io.modelcontextprotocol.server.transport.WebMvcSseServerTransportProvider;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.servlet.function.HandlerFunction;
import org.springframework.web.servlet.function.RouterFunction;
import org.springframework.web.servlet.function.ServerRequest;
import org.springframework.web.servlet.function.ServerResponse;
//import org.springframework.web.reactive.function.server.RouterFunction;
//import org.springframework.web.reactive.function.server.RouterFunctions;
//import org.springframework.web.reactive.function.server.ServerResponse;

import java.util.Map;
import java.util.Optional;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.function.Supplier;

@Slf4j
@Configuration
public class DynamicRouterRegistry{

    private final Map<String, RouterFunction<ServerResponse>> registry = new ConcurrentHashMap<>();
    /**
     * 暴露一个 RouterFunction，将已有的动态路由组合在一起
     */
    @Bean
    public RouterFunction<ServerResponse> mainRouter() {
        return request -> this.findHandler(request).map(Optional::of).orElseGet(Optional::empty);
    }

    /**
     * 对外提供添加路由的方法
     */
    public void register(String routeId, WebMvcSseServerTransportProvider transportProvider) {
        registry.put(routeId, transportProvider.getRouterFunction());
        log.info("Registered route: {}", routeId);
    }

    /** 注销路由 */
    public void unregister(String routeId) {
        registry.remove(routeId);
        log.info("Unregistered route: {}", routeId);
    }

    private Optional<HandlerFunction<ServerResponse>> findHandler(ServerRequest request) {
        for (RouterFunction<ServerResponse> rf : registry.values()) {
            Optional<HandlerFunction<ServerResponse>> candidate = rf.route(request);
            if (candidate.isPresent()) return candidate;
        }
        return Optional.empty();
    }


}
